from collections import deque
import re
import subprocess
from concurrent.futures import ThreadPoolExecutor
import sys
from kea2.hypiumDriver import UiDriver
from enum import Enum

class StampType(Enum):
    ANIMATION_START = "animation_start"
    ANIMATION_END = "animation_end"
    FRAME = "frame"
    
STAMP_MAP = {
    StampType.ANIMATION_START: r".+\s(\d+\.\d+):.+H:Animation start.+",
    StampType.ANIMATION_END: r".+\s(\d+\.\d+):.+H:Animation end.+",
    StampType.FRAME: r".+\s(\d+\.\d+):.+H:Waiting for Present Fence.+"
}


class TraceRecorder:
    
    _instance = None
    
    def __new__(cls, *args, **kwargs):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self, driver: UiDriver):
        self.file_name = "temp.trace"
        self.driver = driver
        self.state = "None"

    def start(self):
        if self.state != "start":
            self.state = "start"
            self.driver.shell(
                "hitrace --trace_begin --record app commercial ace"
            )

    def end(self):
        if self.state != "end":
            self.state = "end"
            with open(self.file_name, "w") as fp:
                subprocess.run(
                    "hdc shell hitrace --trace_finish".split(),
                    stdout=fp
                )

    def get_duration(self, key):
        """Get duration of the target key.
        """
        if self.state != "end":
            print("TraceRecorder is not ended yet.")
            return None

        res = []
        with open(self.file_name, "r") as fp:
            res = [line.strip() for line in fp if key in line]

        if len(res) != 2:
            print(f"No duration for {key}")
            return None

        PATTERN = r".+\s(\d+\.\d+):\s.+"
        m1 = re.match(PATTERN, res[0])
        m2 = re.match(PATTERN, res[1])

        start_time = m1.group(1)
        end_time = m2.group(1)

        return float(end_time) - float(start_time)
    
    
    @classmethod
    def cal_process_node(self):
        """Calculate the process node count.
        """
        ans = dict()
        frame = 0
        with open("/Users/atria/Desktop/coding/hypiumKea/Kea2/out.trace", "r") as fp:
            for line in fp:
                if "H:RenderFrame" in line:
                    frame += 1
                match = re.search(r"ProcessedNodes:\s*(\d+)", line)
                #  RSUniRenderThre-1794    (   1479) [001] .... 511455.541490: tracing_mark_write: B|1479|H:RSUniRenderThread::Render() the number of total ProcessedNodes: 14
                if match:
                    count = int(match.group(1))
                    ans[frame] = ans.get(frame, 0) + count
        return max(ans.values()) if ans else 0
    
    
    @classmethod
    def cal_frame_refresh(self, tag="APP_LIST_FLING"):
        refresh_rate_tag = r".+CommitAndReleaseLayers rate: (\d+),.+"

        def get_refresh_rate(line):
            match = re.match(refresh_rate_tag, line)
            if match:
                return int(match.group(1))
            return -1
        
        def get_stamp(line, stamp_type: StampType):
            match = re.match(STAMP_MAP[stamp_type], line)
            if match:
                return float(match.group(1))
            return -1

        with open("/Users/atria/Desktop/coding/hypiumKea/Kea2/out2.trace", "r") as fp:
            cur_refresh_rate = -1
            last_frame_stamp = -1
            cal_flag = False
            _animation_start = -1
            _animation_end = -1
            for line in fp:
                if tag in line:
                    if (animation_start := get_stamp(line, StampType.ANIMATION_START)) != -1:
                        cal_flag = True
                        frame_lagging = 0
                        _animation_start = animation_start
                    if (animation_end := get_stamp(line, StampType.ANIMATION_END)) != -1:
                        cal_flag = False
                        frame_lagging
                        _animation_end = animation_end
                if not cal_flag:
                    continue
                if (refresh_rate := get_refresh_rate(line)) != -1:
                    if cur_refresh_rate != refresh_rate:
                        cur_refresh_rate = refresh_rate
                        last_frame_stamp = -1  # Reset last frame time when refresh rate changes
                if cur_refresh_rate != -1 and (frame_stamp := get_stamp(line, StampType.FRAME)) != -1:
                    if last_frame_stamp != -1:
                        frame_time = frame_stamp - last_frame_stamp
                        print(f"Frame time: {frame_time:.3f} ms")
                        frame_lagging += max((frame_time - (1 / cur_refresh_rate) * 1.5), 0)
                    last_frame_stamp = frame_stamp
        animation_duration = _animation_end - _animation_start
        assert frame_lagging / animation_duration < 0.005, "卡顿率(即流畅度) = 【每帧丢帧时间累计总和(ms)】/ 总时长(s)，须小于等于5ms/s。"

    def __del__(self):
        import os
        os.remove(self.file_name)


class MemDumper:
    """Dump memory usage of the target key.
    """
    def __init__(self, driver: UiDriver, package: str):
        self.driver = driver
        self.cache = deque()
        self.package = package
        self.thread_pool = ThreadPoolExecutor(max_workers=1)

    def dump(self, package: str = None):
        """Dump memory usage of the target key.
        Args:
            package (str): Package name to dump memory usage.
            If None, use the package name set in the constructor.
        """
        res = self.driver.shell("ps -ef")
        package = package or self.package
        pid = None
        for line in res.splitlines()[1:]:
            if package in line:
                pid = line.split()[1]
                break
        
        if not pid:
            print(f"Error: {package} not found in process list.")
            return
        
        res = self.driver.shell(
            f"hidumper --mem {pid}"
        )
        
        for line in res.splitlines():
            if line.strip().startswith("Total"):
                mem_usage = line.split()[1]
                print(mem_usage)


if __name__ == "__main__":
    # Example usage
    from kea2.hypiumDriver import HDriver
    HDriver.initializer("23E0224307001843")
    driver = HDriver()
    recorder = MemDumper(
        driver,
        package="com.zhihu.hmos"
    )
    recorder.dump()
    recorder.dump()